#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# 接收 /wrench_fake 话题，050 和 000 交替
# 如果接收到的力不超过期望力020，则控制机械臂往初始位置运动

import time
import numpy as np
import rospy
import subprocess
from moveit_commander import MoveGroupCommander
from geometry_msgs.msg import Pose, WrenchStamped,TwistStamped

class ForceControl:
    def __init__(self):
        self.force_torque_filtered = np.zeros((3, 1))
        self.force_torque_reference = np.zeros((3, 1))
        self.force_torque_reference[1] = 0  # 参考力
        self.position_reference = np.zeros((3, 1))
        self.tcp_frame_cur_p = np.zeros((3, 1))
        self.tcp_frame_cur_M = np.eye(3)
        self.velocity = np.zeros((3, 1))
        self.acceleration = np.zeros((3, 1))
        self.K_position = np.eye(3) * 1.1
        self.B_position = np.eye(3) * 100
        self.M_position = np.eye(3) * 280
        self.dt = 1
        self.first_get_coord = True
        self.group = MoveGroupCommander("manipulator")

        # 订阅 /force_data 话题
        self.wrench_sub = rospy.Subscriber('/force_data', TwistStamped, self.force_callback)
        
    def force_callback(self, msg):
        self.force_torque_filtered[0] = msg.twist.linear.z*5
        # self.force_torque_filtered[0] = 0.2*10
        self.force_torque_filtered[1] = 0
        self.force_torque_filtered[2] = 0

    def update_current_pose(self):
        current_pose = self.group.get_current_pose().pose
        self.tcp_frame_cur_p[0, 0] = current_pose.position.x
        self.tcp_frame_cur_p[1, 0] = current_pose.position.y
        self.tcp_frame_cur_p[2, 0] = current_pose.position.z
        if self.first_get_coord:
            self.position_reference[0, 0] = current_pose.position.x
            self.position_reference[1, 0] = current_pose.position.y
            self.position_reference[2, 0] = current_pose.position.z
            self.first_get_coord = False

    def calculate_new_pose(self):
        force_difference = self.force_torque_filtered - self.force_torque_reference
        position_difference = self.tcp_frame_cur_p - self.position_reference
        Kp_position_term = np.dot(self.K_position, position_difference)
        B_velocity_term = np.dot(self.B_position, self.velocity)
        
        self.acceleration = np.dot(np.linalg.inv(self.M_position), 
                                   (force_difference - Kp_position_term - B_velocity_term))
        self.velocity += (self.acceleration * self.dt)
        delta_position = self.velocity * self.dt
        
        new_position = self.tcp_frame_cur_p - delta_position
        # 打印加速度、速度和新位置
        # print("Acceleration:\n", self.acceleration)
        # print("Velocity:\n", self.velocity)
        print("New Position:\n", new_position)
        return new_position
    
    def force_below_threshold(self, threshold):
        return abs(self.force_torque_filtered[0, 0]) <= threshold[0, 0]  # 检查z方向的力
    
def start_force_sensor():
    # 启动力传感器节点
    sensor_process = subprocess.Popen(['python3', '/home/jetson/zc3/src/visual_control/src/force_run3.py'])
    print("力传感器节点打开")
    return sensor_process

def stop_force_sensor(sensor_process):
    # 停止力传感器节点
    sensor_process.terminate()
    print("力传感器节点关闭")


def force_control_main():
    print("\n力控开始……")
    group_name = "manipulator"
    move_group = MoveGroupCommander(group_name)
    force_control = ForceControl()
    
    # 设置初始位置 0/35/45/-80/0
    # initial_pose = Pose()
    # initial_pose.position.x = 0.40764351154554024
    # initial_pose.position.y = 0.00011870430676285974
    # initial_pose.position.z = 0.25493990404932604
    # initial_pose.orientation.w = 1.0
    
    # move_group.set_pose_target(initial_pose)
    # move_group.go(wait=True)
    move_group.set_named_target('middle')
    move_group.go()
    move_group.stop()
    move_group.clear_pose_targets()
    print("到位置")

    # 启动力传感器节点
    # sensor_process = start_force_sensor()
    last_state = None # 初始化 last_state 变量
    
    try:
        while not rospy.is_shutdown():
            force_control.update_current_pose()

            # 检查接收到的力是否超过期望力
            if  force_control.force_torque_filtered[0] == 0:
                # 如果接收到的力不超过期望力，控制机械臂运动到初始位置
                move_group.set_named_target('middle')
                move_group.go(wait=True)
                if last_state != "return_to_initial":
                    print("回到送食位置")
                    last_state = "return_to_initial"
            else:
                new_position = force_control.calculate_new_pose()
                
                pose_target = move_group.get_current_pose().pose
                pose_target.position.x = new_position[0, 0]
                pose_target.position.y = new_position[1, 0]
                pose_target.position.z = new_position[2, 0]
                pose_target.orientation.x = 0
                pose_target.orientation.y = 0
                pose_target.orientation.z = 0
                pose_target.orientation.w = 0
                
                move_group.set_pose_target(pose_target)
                move_group.go(wait=False)
                if last_state != "impedance_control":
                    last_state = "impedance_control"
                time.sleep(1)
            
    except KeyboardInterrupt:
        print("Motion stopped by user")
    finally:
        # stop_force_sensor(sensor_process) # 停止力传感器节点
        move_group.stop()
        move_group.clear_pose_targets()
